# states/global_state.py（全局状态管理，统筹系统所有资源）
from typing import Optional, Dict, Any, List
from datetime import datetime
import time
import threading
import psutil
from loguru import logger

# 导入已验证的核心模块
from utils.logger import Logger
from utils.exceptions import RedisError
from utils.constants import (
    REDIS_HOST, REDIS_PORT, REDIS_DB, REDIS_PASSWORD,
    CACHE_EXPIRE_GLOBAL_STATE,  # 新增全局状态缓存过期常量
    AgentType, TaskState
)
from middleware.redis_adapter import get_redis_adapter
from states.agent_state import AgentState
from states.task_state import TaskStateManager


class GlobalState:
    """全局状态管理器（单例模式）"""
    _instance = None
    _lock = threading.Lock()  # 单例锁
    _update_interval = 10  # 全局状态更新间隔（10秒）

    def __new__(cls):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super(GlobalState, cls).__new__(cls)
                    cls._instance._init()
        return cls._instance

    def _init(self):
        """初始化全局状态"""
        # 依赖组件
        self.redis_client = get_redis_adapter()
        self.logger = Logger.get_logger(agent_name="global_state_manager")

        # 全局状态核心字段（实时更新）
        self.system_metrics: Dict[str, float] = {  # 系统级资源使用率
            "total_cpu_usage": 0.0,
            "total_mem_usage": 0.0,
            "total_disk_usage": 0.0,
            "online_agent_count": 0,  # 在线Agent实例总数
            "healthy_agent_count": 0  # 健康Agent实例总数
        }

        self.agent_cluster_stats: Dict[str, Dict[str, int]] = {  # 各类型Agent集群统计
            # 格式：{agent_type: {"total": 总数, "healthy": 健康数, "loading": 负载高, "unhealthy": 不健康}}
        }

        self.task_global_stats: Dict[str, int] = {  # 全局任务统计
            "total_task_count": 0,
            "pending_task_count": 0,
            "running_task_count": 0,
            "success_task_count": 0,
            "failed_task_count": 0,
            "retrying_task_count": 0
        }

        self.global_config: Dict[str, Any] = {  # 全局配置（可动态更新）
            "max_agent_concurrent": 5,
            "task_max_retries": 3,
            "system_resource_threshold": {  # 系统资源告警阈值
                "cpu": 90.0,
                "mem": 85.0,
                "disk": 95.0
            }
        }

        # 启动定时更新线程（守护线程）
        self._start定时更新()
        self.logger.success("全局状态管理器初始化完成，已启动定时更新线程")

    def _collect_system_metrics(self) -> None:
        """采集系统级资源使用率（全局CPU/内存/磁盘）"""
        try:
            # 全局CPU使用率（10秒内平均值，避免波动）
            self.system_metrics["total_cpu_usage"] = round(psutil.cpu_percent(interval=0.5), 2)

            # 全局内存使用率
            mem_info = psutil.virtual_memory()
            self.system_metrics["total_mem_usage"] = round(mem_info.percent, 2)

            # 全局磁盘使用率（系统根分区，适配Windows/Linux）
            import os
            if os.name == "nt":  # Windows
                disk_partition = "C:"
            else:  # Linux/Mac
                disk_partition = "/"
            disk_info = psutil.disk_usage(disk_partition)
            self.system_metrics["total_disk_usage"] = round(disk_info.percent, 2)

            self.logger.debug(
                f"系统资源采集完成: CPU={self.system_metrics['total_cpu_usage']}% | 内存={self.system_metrics['total_mem_usage']}% | 磁盘={self.system_metrics['total_disk_usage']}%")
        except Exception as e:
            self.logger.error(f"系统资源采集失败: {str(e)}")
            # 采集失败时保持上次值，避免状态突变
            pass

    def _stats_agent_cluster(self) -> None:
        """统计各类型Agent集群状态（总数/健康数/负载数/不健康数）"""
        try:
            agent_types = [agent_type.value for agent_type in AgentType]
            total_online = 0
            total_healthy = 0

            for agent_type in agent_types:
                # 获取该类型所有实例（Redis模糊查询）
                pattern = f"agent:state:{agent_type}:*"
                instance_keys = self.redis_client.client.keys(pattern)
                total = len(instance_keys)

                # 统计该类型各状态实例数
                healthy = 0
                loading = 0
                unhealthy = 0

                for key in instance_keys:
                    key_str = key.decode("utf-8") if isinstance(key, bytes) else key
                    state_data = self.redis_client.hgetall(key_str)
                    status = state_data.get("status", "unhealthy")
                    if status == "healthy":
                        healthy += 1
                    elif status == "loading":
                        loading += 1
                    else:
                        unhealthy += 1

                # 更新该类型统计
                self.agent_cluster_stats[agent_type] = {
                    "total": total,
                    "healthy": healthy,
                    "loading": loading,
                    "unhealthy": unhealthy
                }

                total_online += total
                total_healthy += healthy

            # 更新全局Agent统计
            self.system_metrics["online_agent_count"] = total_online
            self.system_metrics["healthy_agent_count"] = total_healthy

            self.logger.debug(f"Agent集群统计完成: 在线实例数={total_online} | 健康实例数={total_healthy}")
        except Exception as e:
            self.logger.error(f"Agent集群统计失败: {str(e)}")

    def _stats_global_tasks(self) -> None:
        """统计全局任务状态分布（各状态任务数）"""
        try:
            # 重置统计数
            self.task_global_stats = {
                "total_task_count": 0,
                "pending_task_count": 0,
                "running_task_count": 0,
                "success_task_count": 0,
                "failed_task_count": 0,
                "retrying_task_count": 0
            }

            # 模糊查询所有任务状态键
            pattern = "task:state:*"
            task_keys = self.redis_client.client.keys(pattern)
            total_tasks = len(task_keys)
            self.task_global_stats["total_task_count"] = total_tasks

            # 统计各状态任务数
            for key in task_keys:
                key_str = key.decode("utf-8") if isinstance(key, bytes) else key
                state_data = self.redis_client.hgetall(key_str)
                task_state = state_data.get("state", TaskState.PENDING.value)

                if task_state == TaskState.PENDING.value:
                    self.task_global_stats["pending_task_count"] += 1
                elif task_state == TaskState.RUNNING.value:
                    self.task_global_stats["running_task_count"] += 1
                elif task_state == TaskState.SUCCESS.value:
                    self.task_global_stats["success_task_count"] += 1
                elif task_state == TaskState.FAILED.value:
                    self.task_global_stats["failed_task_count"] += 1
                elif task_state == TaskState.RETRYING.value:
                    self.task_global_stats["retrying_task_count"] += 1

            self.logger.debug(
                f"全局任务统计完成: 总任务数={total_tasks} | 运行中={self.task_global_stats['running_task_count']} | 成功={self.task_global_stats['success_task_count']} | 失败={self.task_global_stats['failed_task_count']}")
        except Exception as e:
            self.logger.error(f"全局任务统计失败: {str(e)}")

    def _sync_to_redis(self) -> bool:
        """同步全局状态到Redis（供监控面板/其他模块查询）"""
        try:
            # 1. 系统资源指标（Hash）
            system_key = "global:state:system_metrics"
            self.redis_client.hsetall(system_key, {
                "total_cpu_usage": str(self.system_metrics["total_cpu_usage"]),
                "total_mem_usage": str(self.system_metrics["total_mem_usage"]),
                "total_disk_usage": str(self.system_metrics["total_disk_usage"]),
                "online_agent_count": str(self.system_metrics["online_agent_count"]),
                "healthy_agent_count": str(self.system_metrics["healthy_agent_count"]),
                "updated_at": str(int(datetime.now().timestamp()))
            })
            self.redis_client.expire(system_key, CACHE_EXPIRE_GLOBAL_STATE)

            # 2. Agent集群统计（Hash，按类型存储）
            agent_cluster_key = "global:state:agent_cluster"
            agent_cluster_data = {}
            for agent_type, stats in self.agent_cluster_stats.items():
                for stat_key, value in stats.items():
                    agent_cluster_data[f"{agent_type}:{stat_key}"] = str(value)
            self.redis_client.hsetall(agent_cluster_key, agent_cluster_data)
            self.redis_client.expire(agent_cluster_key, CACHE_EXPIRE_GLOBAL_STATE)

            # 3. 全局任务统计（Hash）
            task_stats_key = "global:state:task_stats"
            self.redis_client.hsetall(task_stats_key, {
                k: str(v) for k, v in self.task_global_stats.items()
            })
            self.redis_client.expire(task_stats_key, CACHE_EXPIRE_GLOBAL_STATE)

            # 4. 全局配置（String，JSON序列化）
            config_key = "global:state:config"
            import json
            self.redis_client.set(config_key, json.dumps(self.global_config, ensure_ascii=False),
                                  ex=CACHE_EXPIRE_GLOBAL_STATE)

            self.logger.debug("全局状态同步到Redis成功")
            return True
        except Exception as e:
            self.logger.error(f"全局状态同步到Redis失败: {str(e)}")
            raise RedisError(message="Global state sync to Redis failed", context={"error": str(e)}) from e

    def _check_resource_alarm(self) -> None:
        """检查系统资源是否超出阈值，触发告警"""
        threshold = self.global_config["system_resource_threshold"]
        alarms = []
        if self.system_metrics["total_cpu_usage"] >= threshold["cpu"]:
            alarms.append(f"CPU使用率超标: {self.system_metrics['total_cpu_usage']}%（阈值{threshold['cpu']}%）")
        if self.system_metrics["total_mem_usage"] >= threshold["mem"]:
            alarms.append(f"内存使用率超标: {self.system_metrics['total_mem_usage']}%（阈值{threshold['mem']}%）")
        if self.system_metrics["total_disk_usage"] >= threshold["disk"]:
            alarms.append(f"磁盘使用率超标: {self.system_metrics['total_disk_usage']}%（阈值{threshold['disk']}%）")

        if alarms:
            alarm_msg = "系统资源告警: " + " | ".join(alarms)
            self.logger.warning(alarm_msg)
            # 这里可扩展：发送告警通知（邮件/短信/企业微信）

    def _update_global_state(self) -> None:
        """更新全局状态（采集+统计+同步+告警）"""
        try:
            self._collect_system_metrics()  # 采集系统资源
            self._stats_agent_cluster()  # 统计Agent集群
            self._stats_global_tasks()  # 统计全局任务
            self._sync_to_redis()  # 同步到Redis
            self._check_resource_alarm()  # 资源告警检查
        except Exception as e:
            self.logger.error(f"全局状态更新失败: {str(e)}")

    def _start定时更新(self) -> None:
        """启动定时更新线程"""

        def update_loop():
            while True:
                self._update_global_state()
                time.sleep(self._update_interval)

        update_thread = threading.Thread(target=update_loop, daemon=True)
        update_thread.start()
        self.logger.info(f"全局状态定时更新线程已启动，更新间隔={self._update_interval}秒")

    # ---------------------- 对外查询接口（供其他模块调用） ----------------------
    def get_system_metrics(self) -> Dict[str, float]:
        """获取系统资源指标"""
        return self.system_metrics.copy()

    def get_agent_cluster_stats(self, agent_type: Optional[str] = None) -> Dict[str, Any]:
        """获取Agent集群统计（指定类型或所有类型）"""
        if agent_type:
            return self.agent_cluster_stats.get(agent_type, {"total": 0, "healthy": 0, "loading": 0, "unhealthy": 0})
        return self.agent_cluster_stats.copy()

    def get_task_global_stats(self) -> Dict[str, int]:
        """获取全局任务统计"""
        return self.task_global_stats.copy()

    def get_global_config(self) -> Dict[str, Any]:
        """获取全局配置"""
        return self.global_config.copy()

    # ---------------------- 对外更新接口（动态调整配置） ----------------------
    def update_global_config(self, config_updates: Dict[str, Any]) -> bool:
        """更新全局配置（支持部分更新）"""
        try:
            with self._lock:
                self.global_config.update(config_updates)
                # 同步到Redis
                config_key = "global:state:config"
                import json
                self.redis_client.set(config_key, json.dumps(self.global_config, ensure_ascii=False),
                                      ex=CACHE_EXPIRE_GLOBAL_STATE)
            self.logger.success(f"全局配置更新成功: {config_updates}")
            return True
        except Exception as e:
            self.logger.error(f"全局配置更新失败: {str(e)}")
            return False


# ===================== 补充 constants.py 缺失常量 =====================
# 打开 utils/constants.py，添加全局状态缓存过期时间（已在上面代码中引用）
"""
# 在 utils/constants.py 中添加
CACHE_EXPIRE_GLOBAL_STATE: Final[int] = 3600  # 全局状态缓存1小时
"""

# ===================== 使用示例 =====================
if __name__ == "__main__":
    import sys

    sys.path.append("..")

    # 初始化全局状态管理器（单例）
    global_state = GlobalState()

    # 等待1秒，让第一次更新完成
    time.sleep(1)

    # 1. 查询系统资源指标
    print("=== 系统资源指标 ===")
    system_metrics = global_state.get_system_metrics()
    for key, value in system_metrics.items():
        print(f"{key}: {value}")

    # 2. 查询Agent集群统计
    print("\n=== Agent集群统计 ===")
    agent_stats = global_state.get_agent_cluster_stats()
    for agent_type, stats in agent_stats.items():
        print(
            f"{agent_type}: 总数={stats['total']} | 健康={stats['healthy']} | 负载高={stats['loading']} | 不健康={stats['unhealthy']}")

    # 3. 查询全局任务统计
    print("\n=== 全局任务统计 ===")
    task_stats = global_state.get_task_global_stats()
    for key, value in task_stats.items():
        print(f"{key}: {value}")

    # 4. 更新全局配置（示例：调整CPU告警阈值）
    print("\n=== 更新全局配置 ===")
    update_result = global_state.update_global_config({"system_resource_threshold": {"cpu": 85.0}})
    print(f"配置更新结果: {'成功' if update_result else '失败'}")
    print("更新后配置:", global_state.get_global_config()["system_resource_threshold"])

    # 保持运行，观察定时更新
    print("\n=== 全局状态管理器持续运行中（按Ctrl+C退出）===")
    try:
        while True:
            time.sleep(10)
    except KeyboardInterrupt:
        print("\n退出全局状态管理器")